package com.tjbklx1.zkclient.lock;

import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;

import org.I0Itec.zkclient.IZkDataListener;
import org.I0Itec.zkclient.ZkClient;
import org.I0Itec.zkclient.exception.ZkNoNodeException;

public class BaseDistributedLock {

	private final ZkClientExt client;
	private final String path;
	private final String basePath;
	private final String lockName;
	private static final Integer MAX_RETRY_COUNT = 10;

	public BaseDistributedLock(ZkClientExt client, String path, String lockName) {
		this.client = client;
		this.basePath = path;
		this.path = path.concat("/").concat(lockName);
		this.lockName = lockName;
	}

	/**
	 * 删除节点
	 * 
	 * @param ourPath
	 * @throws Exception
	 */
	private void deleteOurPath(String ourPath) throws Exception {
		client.delete(ourPath);
	}

	/**
	 * 创建临时节点
	 * 
	 * @param client
	 * @param path
	 * @return
	 * @throws Exception
	 */
	private String createLockNode(ZkClient client, String path) throws Exception {
		return client.createEphemeralSequential(path, null);
	}

	/**
	 * 等待获得锁 (核心方法)
	 * 
	 * @param startMillis
	 * @param millisToWait
	 * @param ourPath
	 * @return
	 * @throws Exception
	 */
	private boolean waitToLock(long startMillis, Long millisToWait, String ourPath) throws Exception {

		boolean haveTheLock = false;
		boolean doDelete = false;

		try {
			while (!haveTheLock) {
				// 获得子节点
				List<String> children = getSortedChildren();
				// 节点的名称
				String sequenceNodeName = ourPath.substring(basePath.length() + 1);
				// 查找节点
				int ourIndex = children.indexOf(sequenceNodeName);
				if (ourIndex < 0) {
					throw new ZkNoNodeException("节点没有找到: " + sequenceNodeName);
				}
				// 是否是第一个
				boolean isGetTheLock = ourIndex == 0;
				// 前一个节点
				String pathToWatch = isGetTheLock ? null : children.get(ourIndex - 1);

				// 是第一个就获得锁 ,返回TRUE,否则 等待
				if (isGetTheLock) {
					haveTheLock = true;
				} else {
					// 前一个节点的名称
					String previousSequencePath = basePath.concat("/").concat(pathToWatch);
					final CountDownLatch latch = new CountDownLatch(1);
					// 监听删除事件
					final IZkDataListener previousListener = new IZkDataListener() {
						public void handleDataDeleted(String dataPath) throws Exception {
							latch.countDown();
						}

						public void handleDataChange(String dataPath, Object data) throws Exception {
							// ignore
						}
					};

					try {
						// 如果节点不存在会出现异常
						// 订阅上一个节点的删除事件
						client.subscribeDataChanges(previousSequencePath, previousListener);

						if (millisToWait != null) {
							millisToWait -= (System.currentTimeMillis() - startMillis);
							startMillis = System.currentTimeMillis();
							if (millisToWait <= 0) {
								doDelete = true; // timed out - delete our node
								break;
							}

							latch.await(millisToWait, TimeUnit.MICROSECONDS);
						} else {
							latch.await();
						}
					} catch (ZkNoNodeException e) {
						// ignore
					} finally {
						client.unsubscribeDataChanges(previousSequencePath, previousListener);
					}
				}
			}
		} catch (Exception e) {
			// 发生异常需要删除节点
			doDelete = true;
			throw e;
		} finally {
			// 如果需要删除节点
			if (doDelete) {
				deleteOurPath(ourPath);
			}
		}
		return haveTheLock;
	}

	/**
	 * 获取LOCK number
	 * 
	 * @param str
	 * @param lockName
	 * @return
	 */
	private String getLockNodeNumber(String str, String lockName) {
		int index = str.lastIndexOf(lockName);
		if (index >= 0) {
			index += lockName.length();
			return index <= str.length() ? str.substring(index) : "";
		}
		return str;
	}

	/**
	 * 获取子节点信息,并排序
	 * 
	 * @return
	 * @throws Exception
	 */
	private List<String> getSortedChildren() throws Exception {
		try {
			List<String> children = client.getChildren(basePath);
			Collections.sort(children, new Comparator<String>() {
				public int compare(String lhs, String rhs) {
					return getLockNodeNumber(lhs, lockName).compareTo(getLockNodeNumber(rhs, lockName));
				}
			});
			return children;
		} catch (ZkNoNodeException e) {
			client.createPersistent(basePath, true);
			return getSortedChildren();
		}
	}

	/**
	 * 释放节点
	 * 
	 * @param lockPath
	 * @throws Exception
	 */
	protected void releaseLock(String lockPath) throws Exception {
		deleteOurPath(lockPath);

	}

	/**
	 * 试图创建锁
	 * 
	 * @param time
	 * @param unit
	 * @return
	 * @throws Exception
	 */
	protected String attemptLock(long time, TimeUnit unit) throws Exception {

		final long startMillis = System.currentTimeMillis();
		final Long millisToWait = (unit != null) ? unit.toMillis(time) : null;

		String ourPath = null;
		boolean hasTheLock = false;
		boolean isDone = false;
		int retryCount = 0;

		// 网络闪断需要重试一试
		while (!isDone) {
			isDone = true;

			try {
				// 创建临时序列节点
				ourPath = createLockNode(client, path);
				// 等待获得锁
				hasTheLock = waitToLock(startMillis, millisToWait, ourPath);
			} catch (ZkNoNodeException e) {
				// 失败了,多次重试,超过重试次数 返回异常
				if (retryCount++ < MAX_RETRY_COUNT) {
					isDone = false;
				} else {
					throw e;
				}
			}
		}
		// 如果获得锁就返回 ,否则返回NULL
		if (hasTheLock) {
			return ourPath;
		}
		return null;
	}

}
